"""
按照论文[1] Wang Gangyi, Li Jian, Wei Xinguo. Star identification based on hash map[J]. IEEE Sensors Journal, 2018, 18(4): 1591-1599
和[2] Christian J. A. , Derksen H. , Watkins R. .Lunar crater identification in digital Images[J/OL].J. Astronaut. Sci.,2021,68(4):1056-1144
"""

import numpy as np
from .base import TriadHash
from itertools import product
from utils.ellipse import radius_ellipse


class ChristianHash(TriadHash):
    def __init__(self, catalog_path, **kwargs):
        super().__init__(catalog_path, **kwargs)
        assert "N_b" in self.constants.keys()
        assert "dist" in self.constants.keys()
        self.N_b = self.constants["N_b"]
        self.max_dist = self.constants["dist"][1]
        self.min_dist = self.constants["dist"][0]

    def triad_descriptor(
        self, i, j, k, Q1, Q2, Q3, *args, th=3, **kwargs
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """从陨石坑的参数中计算描述子，C1和C2可以是含有不确定度的椭圆参数"""
        N = Q3.shape[0]
        if len(Q1.shape) == 2:
            Q1 = Q1[None].repeat(N, axis=0)
        if len(Q2.shape) == 2:
            Q2 = Q2[None].repeat(N, axis=0)
        assert Q1.shape[0] == Q2.shape[0] == Q3.shape[0]
        d1 = np.array(radius_ellipse(Q1)).mean(axis=0)
        d2 = np.array(radius_ellipse(Q2)).mean(axis=0)
        d3 = np.array(radius_ellipse(Q3)).mean(axis=0)
        # 按直径从小到大排序
        index = np.argsort((d1, d2, d3), axis=0)
        Q = np.array((Q1, Q2, Q3))[index, np.arange(N)]
        # 计算不变量
        I1 = np.einsum("nii->n", Q[2] @ np.linalg.inv(Q[1])) * np.einsum(
            "nii->n", Q[1] @ np.linalg.inv(Q[2])
        )
        I2 = np.einsum("nii->n", Q[0] @ np.linalg.inv(Q[2])) * np.einsum(
            "nii->n", Q[2] @ np.linalg.inv(Q[0])
        )
        I3 = np.einsum("nii->n", Q[1] @ np.linalg.inv(Q[0])) * np.einsum(
            "nii->n", Q[0] @ np.linalg.inv(Q[1])
        )
        # 量化
        qs = np.log(np.abs((I1, I2, I3)))
        qs = np.int32(
            np.power((qs - self.min_dist) / (self.max_dist - self.min_dist), 3)
            * (self.N_b - 1)
        )
        # 根据误差阈值产生不同正负误差组合方式
        th = np.array(list(product(np.arange(-th, th + 1), repeat=3))).T
        # 合并为Hash值
        qs = (
            (qs[0] + th[0, :, None]) * self.N_b * self.N_b
            + (qs[1] + th[1, :, None]) * self.N_b
            + (qs[2] + th[2, :, None])
        )
        # 输出计算所得哈希值
        ijk = np.vstack((i, j, k))
        return qs, ijk, np.argsort(index, axis=0), np.ones_like(d1, dtype=bool)
